package com.liuss.springchatai.controller;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @Author: lss
 * @Date: 2025-04-15-10:15
 * @Description:
 */
@RestController
@RequestMapping("/rag")
@Slf4j
public class RAGInPutController {

    @Resource
    private VectorStore store;

    /**
     * 小文本输入,建议,更准确的回答内容
     * @return
     * @throws IOException
     */
    @GetMapping("/writer")
    public String writer() throws IOException {
        StringBuffer text = new StringBuffer();
        ClassLoader classLoader = getClass().getClassLoader();
        InputStream inputStream = classLoader.getResourceAsStream("斗破苍穹(1章).txt");
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            while ((line = reader.readLine()) != null) {
                text.append(line);
            }
        }
        store.write(Arrays.stream(text.toString().split("\n")).map(Document::new).toList());
        return "success";
    }

    /**
     * 大文本分段输入,不建议 会影响数据的输出
     * @return
     * @throws IOException
     */
    @GetMapping("/writerMaxTxt")
    public String writerMaxText() throws IOException {
        StringBuilder text = new StringBuilder();
        ClassLoader classLoader = getClass().getClassLoader();
        InputStream inputStream = classLoader.getResourceAsStream("斗破苍穹(1-500章).txt");
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            while ((line = reader.readLine()) != null) {
                text.append(line).append("\n");
            }
        }

        List<Document> allDocuments = new ArrayList<>();
        String fullText = text.toString();
        String[] paragraphs = fullText.split("\n");

        StringBuilder currentDocumentContent = new StringBuilder();
        for (String paragraph : paragraphs) {
            if (currentDocumentContent.length() + paragraph.length() + 1 > 2048) {
                allDocuments.add(new Document(currentDocumentContent.toString().trim()));
                currentDocumentContent.setLength(0);
            }
            currentDocumentContent.append(paragraph).append("\n");
        }

        if (currentDocumentContent.length() > 0) {
            allDocuments.add(new Document(currentDocumentContent.toString().trim()));
        }

        store.write(allDocuments);
        log.info("大文本上传完毕.");
        return "success";
    }

    @GetMapping("/writerPdf")
    public String writerPdf() throws IOException {
        //读取pdf文件存入向量库
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader("classpath:/templates/高等学校实验室安全检查项目表（2024年）.pdf",
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());
        store.write(pagePdfDocumentReader.read());
        return "success";
    }

    /**
     * Markdown 文件
     */
    public List<Document> md() {
        MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader("classpath:/lists.md");
        return markdownDocumentReader.read();
    }

}
